'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 6_self_define_document.py
* @Time: 2025/10/29
* @All Rights Reserve By Brtc
'''
from typing import Iterator, AsyncIterator

from langchain_core.document_loaders import BaseLoader
from langchain_core.documents import Document


class CustomDocument(BaseLoader):
    """自定义文档加载器"""
    def __init__(self, file_path):
        self._file_path = file_path

    def lazy_load(self) -> Iterator[Document]:
        """实现懒加载"""
        #逐行读取文件的数据并使用yield返回文档
        with open(self._file_path, "r", encoding="utf-8") as f:
            line_nums = 0
            for line in f:
                yield Document(
                    page_content=line,
                    metadata={"line_number": line_nums, "source":self._file_path},
                )
                line_nums += 1

    async def alazy_load(self) -> AsyncIterator[Document]:
        """异步懒加载"""
        import  aiofiles
        async with aiofiles.open(self._file_path, "r", encoding="utf-8") as f:
            line_nums = 0
            async for line in f:
                yield Document(
                    page_content=line,
                    metadata={"line_number": line_nums, "source":self._file_path},
                )
                line_nums += 1

loader = CustomDocument("./test.txt")
docs = loader.load()

print(docs)
print(len(docs))
print(docs[0].metadata)